/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2024-2024. All rights reserved.
 * Description: 集合通信op_base模块
 * Author: baochangjin
 * Create: 2024-4-2
 */

#include "op_base_rpc.h"
#include <algorithm>
#include <future>
#include <map>
#include <string>
#include <hccl/hccl_types.h>

#include "hccl/base.h"
#include "hccl/hccl_ex.h"
#include "workflow_pub.h"
#include "param_check_pub.h"
#include "rank_consistentcy_checker.h"
#include "dlra_function.h"
#include "dltdt_function.h"
#include "externalinput_pub.h"
#include "../common/src/topo/topoinfo_detect.h"
#include "mem_host_pub.h"
#include "alg_template_base_pub.h"
#include "profiling_manager.h"
#include "adapter_prof.h"
#include "hcom_pub.h"
#include "adapter_qos.h"
#include "dlqos_function.h"
#include "device_capacity.h"
#include "hcom_rpc.h"
#include "param_check_pub.h"
#include "transport_heterog_def.h"

using namespace std;
using namespace hccl;

u64 HcclGetLookupUpdateWorkspace(s32 count, s32 valueDim, HcclDataType keyType, HcclDataType valueType, s32 flags)
{
    u64 opMemSize = 0;

    // +1是为table id申请的内存
    s32 valueItemSize = valueDim * SIZE_TABLE[valueType];
    s64 keyMemSize = (((count + 1) * SIZE_TABLE[keyType] / IPC_MEM_ALIGNMENT_BYTE) + 1) * IPC_MEM_ALIGNMENT_BYTE;
    // 为key的中转申请2倍内存,1份用作shard多线程去重，1份用作send key buf
    keyMemSize += keyMemSize;
    s64 valueMemSize = ((count * valueItemSize / IPC_MEM_ALIGNMENT_BYTE) + 1) * IPC_MEM_ALIGNMENT_BYTE;
    opMemSize = keyMemSize + valueMemSize;

    if (flags == ES_FLAGS_ENABLE_COUNTER) {
        u64 counterSize = Align<u64>(static_cast<u64>(count) * ES_KEY_COUNTER_MEM_BYTES_SIZE, IPC_MEM_ALIGNMENT_BYTE);
        opMemSize += counterSize;
    }

    u64 rdmaEnvelopeMemSize = Align<u64>(sizeof(HcclEsRdmaInfo) * ES_MAX_PS_NUM, IPC_MEM_ALIGNMENT_BYTE);
    opMemSize += rdmaEnvelopeMemSize;

    HCCL_INFO("keyMemSize[%lld], valueMemSize[%lld], rdmaEnvelopeMemSize[%llu], opMemSize[%llu]",
        keyMemSize, valueMemSize, rdmaEnvelopeMemSize, opMemSize);

    return opMemSize;
}

HcclResult HcclRemoteLookup(s32 tag, s32 insertFlag, void *stream, RemoteLookupPubParams &lookupPubInfo)
{
    CHK_PTR_NULL(lookupPubInfo.keys);
    CHK_PTR_NULL(lookupPubInfo.value);
    CHK_PTR_NULL(stream);

    s32 intZerocpyFlag = 0;
    s32 outZerocpyFlag = 0;

    std::shared_ptr<hccl::HcclCommBase> hcclComm = nullptr;
    CHK_RET(HcpuGetCommByGroup(nullptr, hcclComm));

    bool isBoardVersion{};
    CHK_RET((reinterpret_cast<hccl::HcclCommRpc *>(hcclComm.get()))->WorkerGetBoardVersion(isBoardVersion));
    const string tagStr = to_string(tag);

    HcomRemoteOperationParams params{};
    params.tableId = lookupPubInfo.tableId;
    params.keyAddr = lookupPubInfo.keys;
    params.value = lookupPubInfo.value;
    params.count = lookupPubInfo.keyMaxNum;
    params.valueDim = lookupPubInfo.valueItemSize / SIZE_TABLE[HCCL_DATA_TYPE_FP32];
    params.tag = tag;
    params.stream = stream;
    params.insertOption = insertFlag;
    params.flags = 0;
    params.intZerocpyFlag = intZerocpyFlag;
    params.outZerocpyFlag = outZerocpyFlag;
    params.disableUnique = true;
    params.keyType = HCCL_DATA_TYPE_UINT64;
    params.valueType = HCCL_DATA_TYPE_FP32;

    u64 memSize = HcclGetLookupUpdateWorkspace(params.count, params.valueDim, params.keyType, params.valueType,
        params.flags);

    CHK_RET((reinterpret_cast<hccl::HcclCommRpc *>(hcclComm.get()))
        ->CreateOrUpdateRemoteOpBasedResources(memSize, tagStr));

    if (isBoardVersion) {
        CHK_RET((reinterpret_cast<hccl::HcclCommRpc *>(hcclComm.get()))->CollRemoteLookUp(params));
    } else {
        HcclRequest request = nullptr;
        CHK_RET((reinterpret_cast<hccl::HcclCommRpc *>(hcclComm.get()))->
            RemoteLookUp(tag, insertFlag, stream, intZerocpyFlag, outZerocpyFlag, lookupPubInfo,
            &request));
        CHK_RET((reinterpret_cast<hccl::HcclCommRpc *>(hcclComm.get()))->
            WaitLookupData(tag, lookupPubInfo.value, stream, outZerocpyFlag, request));
    }

    return HCCL_SUCCESS;
}

HcclResult HcclCollRemoteLookupOrUpdatePairedOrUniqued(const HcomRemoteOperationParams &params, EsOpType opType)
{
    HCCL_RUN_INFO("HcclCollRemoteLookupUniquedPairedOrUniqued start, tag[%d] key count[%d] value dim size[%d]",
        params.tag, params.count, params.valueDim);

    shared_ptr<hccl::HcclCommBase> hcclComm = nullptr;
    const char *group = params.group.empty() ? nullptr : params.group.c_str();
    CHK_RET(HcpuGetCommByGroup(group, hcclComm));

    const string tagStr = to_string(params.tag);

    // 获取memory size
    u64 memSize = HcclGetLookupUpdateWorkspace(params.count, params.valueDim, HCCL_DATA_TYPE_UINT64,
                                               HCCL_DATA_TYPE_FP32, params.flags);
    // 获取workspace资源
    (reinterpret_cast<hccl::HcclCommRpc *>(hcclComm.get()))->CreateOrUpdateRemoteOpBasedResources(memSize, tagStr);
    // 调用单算子
    if (opType == EsOpType::LOOKUP) {
        CHK_RET((reinterpret_cast<hccl::HcclCommRpc *>(hcclComm.get()))->CollRemoteLookupPaired(params));
    } else if (opType == EsOpType::LOOKUP_UNIQUED) {
        CHK_RET((reinterpret_cast<hccl::HcclCommRpc *>(hcclComm.get()))->CollRemoteLookupUniquedPaired(params));
    } else {
        CHK_RET((reinterpret_cast<hccl::HcclCommRpc *>(hcclComm.get()))->CollRemoteUpdatePaired(params));
    }

    HCCL_RUN_INFO("HcclCollRemoteLookupUniquedPairedOrUniqued success");

    return HCCL_SUCCESS;
}

